"""DeepSORT目标跟踪器封装模块"""

import numpy as np
from typing import List, Dict, Tuple, Optional
from deep_sort_realtime.deepsort_tracker import DeepSort
from config import get_config

class DeepSORTTracker:
    """DeepSORT目标跟踪器封装类"""
    
    def __init__(self):
        """初始化DeepSORT跟踪器"""
        # 从配置加载跟踪器参数
        self.nms_max_overlap = get_config('tracking.nms_max_overlap', 1.0)
        self.max_cosine_distance = get_config('tracking.max_cosine_distance', 0.2)
        self.nn_budget = get_config('tracking.nn_budget', None)
        self.max_age = get_config('tracking.max_age', 30)
        self.n_init = get_config('tracking.n_init', 3)
        self.metric_type = get_config('tracking.metric_type', 'cosine')
        self.before_detected = get_config('tracking.before_detected', False)
        
        # 创建DeepSort跟踪器实例
        self.tracker = DeepSort(
            max_age=self.max_age,
            n_init=self.n_init,
            nms_max_overlap=self.nms_max_overlap,
            max_cosine_distance=self.max_cosine_distance,
            nn_budget=self.nn_budget,
            override_track_class=None,
            embedder='mobilenet',  # 使用轻量级的嵌入模型
            half=True,  # 使用半精度以提高性能
            bgr=True,  # OpenCV图像通常是BGR格式
            embedder_gpu=True  # 如果可用，使用GPU进行嵌入计算
        )
        
        # 跟踪历史记录
        self.track_history = {}
        self.max_history_length = get_config('tracking.max_history_length', 100)
        
        print("DeepSORT跟踪器已初始化")
    
    def update(self, detections: List[Dict], frame: np.ndarray) -> List[Dict]:
        """更新跟踪器
        
        Args:
            detections: 检测结果列表，每个元素包含{x1, y1, x2, y2, confidence, class_id, class_name}
            frame: 当前视频帧
            
        Returns:
            跟踪结果列表，每个元素包含{x1, y1, x2, y2, track_id, confidence, class_id, class_name}
        """
        try:
            # 转换检测结果格式以适应DeepSort
            # DeepSort期望的格式: [(x1, y1, x2, y2, confidence), ...]
            detection_tuples = []
            detection_classes = []
            
            for det in detections:
                # 计算边界框宽度和高度
                x1, y1, x2, y2 = det['x1'], det['y1'], det['x2'], det['y2']
                confidence = det['confidence']
                
                # 添加到检测元组列表
                detection_tuples.append(([x1, y1, x2-x1, y2-y1], confidence, det['class_id']))
                detection_classes.append(det)
            
            # 更新跟踪器
            tracks = self.tracker.update_tracks(detection_tuples, frame=frame)
            
            # 处理跟踪结果
            tracking_results = []
            for track in tracks:
                # 只处理已确认的跟踪
                if not track.is_confirmed():
                    continue
                
                # 获取跟踪ID和边界框
                track_id = track.track_id
                ltrb = track.to_ltrb()  # 格式: [left, top, right, bottom]
                
                # 查找对应的检测结果以获取类别信息
                detection_idx = track.detection_idx
                class_info = None
                if detection_idx is not None and detection_idx < len(detection_classes):
                    class_info = detection_classes[detection_idx]
                
                # 构建跟踪结果
                track_result = {
                    'x1': int(ltrb[0]),
                    'y1': int(ltrb[1]),
                    'x2': int(ltrb[2]),
                    'y2': int(ltrb[3]),
                    'track_id': track_id,
                    'confidence': class_info['confidence'] if class_info else 0.0,
                    'class_id': class_info['class_id'] if class_info else 0,
                    'class_name': class_info['class_name'] if class_info else 'unknown'
                }
                
                tracking_results.append(track_result)
                
                # 更新跟踪历史
                self._update_track_history(track_id, track_result)
            
            return tracking_results
            
        except Exception as e:
            print(f"跟踪更新过程中发生错误: {e}")
            return []
    
    def _update_track_history(self, track_id: int, track_result: Dict) -> None:
        """更新跟踪历史记录
        
        Args:
            track_id: 跟踪ID
            track_result: 当前跟踪结果
        """
        # 如果跟踪ID不存在于历史记录中，创建一个新的历史记录
        if track_id not in self.track_history:
            self.track_history[track_id] = []
        
        # 添加当前位置到历史记录
        self.track_history[track_id].append({
            'x': (track_result['x1'] + track_result['x2']) // 2,  # 中心点x坐标
            'y': (track_result['y1'] + track_result['y2']) // 2,  # 中心点y坐标
            'timestamp': None  # 可以添加时间戳
        })
        
        # 限制历史记录长度
        if len(self.track_history[track_id]) > self.max_history_length:
            self.track_history[track_id].pop(0)
    
    def get_track_history(self, track_id: int) -> List[Dict]:
        """获取指定跟踪ID的历史记录
        
        Args:
            track_id: 跟踪ID
            
        Returns:
            历史位置记录列表
        """
        return self.track_history.get(track_id, [])
    
    def get_all_track_ids(self) -> List[int]:
        """获取所有活动的跟踪ID
        
        Returns:
            跟踪ID列表
        """
        return list(self.track_history.keys())
    
    def remove_track(self, track_id: int) -> None:
        """移除指定的跟踪
        
        Args:
            track_id: 跟踪ID
        """
        if track_id in self.track_history:
            del self.track_history[track_id]
    
    def clear_tracks(self) -> None:
        """清除所有跟踪历史"""
        self.track_history.clear()
        self.tracker = DeepSort(
            max_age=self.max_age,
            n_init=self.n_init,
            nms_max_overlap=self.nms_max_overlap,
            max_cosine_distance=self.max_cosine_distance,
            nn_budget=self.nn_budget,
            override_track_class=None,
            embedder='mobilenet',
            half=True,
            bgr=True,
            embedder_gpu=True
        )
    
    def get_tracking_statistics(self) -> Dict:
        """获取跟踪统计信息
        
        Returns:
            统计信息字典
        """
        return {
            'active_tracks': len(self.track_history),
            'max_history_length': self.max_history_length,
            'max_age': self.max_age,
            'n_init': self.n_init,
            'max_cosine_distance': self.max_cosine_distance
        }
    
    def draw_tracking_results(self, frame: np.ndarray, tracking_results: List[Dict], 
                            show_history: bool = False) -> np.ndarray:
        """在图像上绘制跟踪结果
        
        Args:
            frame: 输入图像
            tracking_results: 跟踪结果列表
            show_history: 是否显示跟踪历史
            
        Returns:
            绘制后的图像
        """
        import cv2
        
        # 创建图像副本
        result_frame = frame.copy()
        
        # 定义颜色映射，用于区分不同的跟踪ID
        color_map = {}
        
        # 绘制每个跟踪结果
        for track in tracking_results:
            track_id = track['track_id']
            
            # 为跟踪ID生成唯一颜色
            if track_id not in color_map:
                # 使用跟踪ID生成伪随机颜色
                color_map[track_id] = ((track_id * 123) % 255, 
                                     (track_id * 321) % 255, 
                                     (track_id * 213) % 255)
            
            color = color_map[track_id]
            x1, y1, x2, y2 = track['x1'], track['y1'], track['x2'], track['y2']
            class_name = track['class_name']
            
            # 绘制边界框
            cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 2)
            
            # 绘制标签
            label = f"ID: {track_id} ({class_name})"
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            label_y = max(y1 - 10, label_size[1])
            
            # 绘制标签背景
            cv2.rectangle(
                result_frame,
                (x1, label_y - label_size[1] - 5),
                (x1 + label_size[0] + 10, label_y + 5),
                color,
                -1
            )
            
            # 绘制标签文本
            cv2.putText(
                result_frame,
                label,
                (x1 + 5, label_y),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 0, 0),
                2
            )
            
            # 绘制跟踪历史
            if show_history:
                history = self.get_track_history(track_id)
                if len(history) > 1:
                    for i in range(1, len(history)):
                        cv2.line(
                            result_frame,
                            (history[i-1]['x'], history[i-1]['y']),
                            (history[i]['x'], history[i]['y']),
                            color,
                            2
                        )
        
        return result_frame
    
    def set_parameters(self, **kwargs) -> None:
        """动态设置跟踪器参数
        
        Args:
            **kwargs: 跟踪器参数
        """
        # 更新配置参数
        if 'max_age' in kwargs:
            self.max_age = kwargs['max_age']
        if 'n_init' in kwargs:
            self.n_init = kwargs['n_init']
        if 'max_cosine_distance' in kwargs:
            self.max_cosine_distance = kwargs['max_cosine_distance']
        if 'nms_max_overlap' in kwargs:
            self.nms_max_overlap = kwargs['nms_max_overlap']
        if 'max_history_length' in kwargs:
            self.max_history_length = kwargs['max_history_length']
        
        # 重新初始化跟踪器
        self.clear_tracks()
        
        print(f"跟踪器参数已更新: {kwargs}")